'''
Created on 10 nov 2009

@author: fredrik
'''
import pygame

def groupcollide_grid(groupa, bgrid, dokilla, dokillb):
    '''
    Tweaked sprite collide function that uses a grid to speed up collision detection.
    Roughly x8 times faster than pygame.sprite.groupcollide in my tests :-)
    '''
    crashed = {}
    i = 0
    width = bgrid.width
    height = bgrid.height
    for sprite in groupa.sprites():
        r = sprite.rect
        spritecollide = r.colliderect
        _dokilla = dokilla
        x_min,y_min = max(0,(r.left/width)-1),max(0,(r.top/height)-1)
        x_max,y_max = max(0,(r.right/width)-1)+1,max(0,(r.bottom/height)-1)+1
        for gy in range(y_min,y_max):
            for gx in range(x_min,x_max):
                if not bgrid[gy][gx]:
                    continue
                for b in bgrid[gy][gx]:
                    i += 1
                    if spritecollide(b.rect):
                        if _dokilla:
                            sprite.kill()
                            _dokilla = False
                        if dokillb:
                            b.kill()
                        crashed.setdefault(sprite,set()).add(b)
        
#    if len(crashed):
#        print "found:%d in %d loops"%(len(crashed),i)
        
    return crashed

class Grid(list):
    '''
    A grid used to hold rough positions of sprites.
    Used by groupcollide_grid as second group.
    @sprites : a list of sprites.
    @size : The pixel size of the world map area (width,height)
            where that are valid for sprites.
    @grid_tiles : desired grid size in tiles.
    '''
    def __init__(self,sprites,size,grid_tiles = 30):
        self.width = width = (size[0]/grid_tiles)+1
        self.height = height = (size[1]/grid_tiles)+1
        self.size = size
        for y in xrange(grid_tiles):
            row = []
            for x in xrange(grid_tiles):
                row.append(None)
            self.append(row)
            
        for s in sprites:
            r = s.rect
            x_min,y_min = max(0,(r.left/width)-1),max(0,(r.top/height)-1)
            x_max,y_max = max(0,(r.right/width)-1)+1,max(0,(r.bottom/height)-1)+1
            try:
                for y in range(y_min,y_max):
                    for x in range(x_min,x_max):
                        if self[y][x] == None:
                            self[y][x] = [s]
                        else:
                            self[y][x].append(s)
            except IndexError,e:
                print "sprite outside map",e

#Test classes
class Block(pygame.sprite.Sprite):
    """moves a clenched fist on the screen, following the mouse"""
    def __init__(self,rect,color=(200,200,200)):
        pygame.sprite.Sprite.__init__(self) #call Sprite initializer
        surf = pygame.image.load('data/bmp/bomb_shot.bmp')
        surf.fill(color)
        surf = pygame.transform.scale(surf, rect[2:])        
        self.image = surf
        self.rect = surf.get_rect().move(rect[:2])
        
    def __repr__(self):
        return '%s'%(self.rect[:4],)

    def update(self):
        "move the fist based on the mouse position"
        pass
                

def print_grid(grid):
    for row in grid:
        print ''.join([str(len(item)) if item else '0' for item in row])

if __name__ == '__main__':
    from pygame.locals import *
    group_a = pygame.sprite.Group()
    group_b = pygame.sprite.Group()
    # (left, top, width, height)
    import random
    rand = random.randint
    def get_rand_rect():
        w = rand(0,20)
        h = rand(0,20)
        return (rand(0,480-w),rand(0,480-h),w,h)
    for r in range(200):
        group_a.add(Block(get_rand_rect()))
    for r in range(300):
        group_b.add(Block(get_rand_rect()))
    
    pygame.init()
    screen = pygame.display.set_mode((640, 480))
    pygame.display.set_caption('Collision test')
#    pygame.mouse.set_visible(0)

#Create The Backgound
    background = pygame.Surface(screen.get_size())
    background = background.convert()
    background.fill((250, 250, 250))

#Display The Background
    screen.blit(background, (0, 0))
    pygame.display.flip()
    loops = 10
#Prepare Game Objects
    clock = pygame.time.Clock()
    allsprites = pygame.sprite.RenderPlain(group_a.sprites())
    allsprites.add(group_b.sprites())
    bgrid = Grid(group_b.sprites(),(480,480))
    print_grid(bgrid)

    import time
    stime = time.time()
    for i in range(loops):
        # Handle ship-shot collisions
        collisions = pygame.sprite.groupcollide(group_a, group_b, False,False)
    etime = time.time()
    print "orig time", etime - stime

    stime = time.time()
    for i in range(loops):
        # Handle ship-shot collisions
        bgrid = Grid(group_b.sprites(),(480,480))
        my_collisions = groupcollide_grid(group_a, bgrid, False,False,)
    print my_collisions
    etime = time.time()
    print "my time", etime - stime
    print "loops with normal collide",len(group_a)*len(group_b)
    if set(collisions.keys()) !=  set(my_collisions.keys()):
        print "Wrong collisions!"
        print sorted(set(collisions.keys()))
        print sorted(set(my_collisions.keys()))
        print len(collisions.keys())
        print len(my_collisions.keys())
#        pygame.quit()
#        raise SystemExit()

    for a,bl in my_collisions.items():
        a.image.fill((0,0,0))
        for b in bl:
            b.image.fill((200,0,255))
    
#Main Loop
    going = True
    while going:
        clock.tick(60)

        #Handle Input Events
        for event in pygame.event.get():
            if event.type == QUIT:
                going = False
            elif event.type == KEYDOWN and event.key == K_ESCAPE:
                going = False

        allsprites.update()

        #Draw Everything
        screen.blit(background, (0, 0))
        allsprites.draw(screen)
        pygame.display.flip()

    pygame.quit()
